import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
import pdb

class ElementWiseOpWrappper(object):

    def __init__(self, opName):
        self.__opName = opName

    def build_model(self):
        X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [3, 4])
        
        # Create one output (ValueInfoProto)
        Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 4])
        
        # Create a node (NodeProto) - This is based on Pad-11
        nodeDef = helper.make_node(
            self.__opName,
            ['X'],                      # inputs
            ['Y'],                      # outputs
            name=f'{self.__opName}_0',
        )
        
        # Create the graph (GraphProto)
        graphDef = helper.make_graph(
            [nodeDef],         # nodes
            'test-model',      # name
            [X],  # inputs
            [Y],               # outputs
        )

        # Create the model (ModelProto)
        return helper.make_model(graphDef, producer_name='ts_quant')
